from dataclasses import dataclass, field
from typing import Any, Tuple, Dict
import torch.nn as nn
import traceback

@dataclass
class Config:
    env_id = "ale/icehockey-v5"
    pixel: bool = True
    gamma = 0.99
    batch_size = 16
    chunk_size = 64
    model_learning_rate = 1e-4
    model_lr_eps = 1e-8
    grad_clip = 1000
    weight_decay = 0.0
    actor_learning_rate = 4e-5
    value_learning_rate = 1e-5
    loss_scale: Dict = field(default_factory=lambda:{'kl':0.1, 'reward':1.0, 'discount':5.0})
    kl_free = 1.0
    dyn_scale = 0.5
    rep_scale = 0.1
    adam_epsilon = 1e-7
    learning_rate_schedule = 0
    replay_size = 1000000 # 重放缓冲区长度，这么长是为了提高稳定性
    replay_initial = 10000 # 重放缓冲区初始化大小
    collect_intervals: int = 5 
    lambda_: float = 0.95
    device = 'cpu'
    discount = 0.997
    discount_lambda = 0.95
    grad_heads = ['decoder', 'reward', 'nonterms']
    gray = False
    stop_reward = 1000000

    test_iters = 100

    precision = 32
    state_size = 30
    hide_size = 200
    belief_size = 200
    reward_step = 1
    embedding_size = 128
    free_nats = float(3.0)
    overshooting_distance = 50

    overshooting_reward_scale = float(0)
    planning_horizon = 10
    grad_clip_norm = 100
    train_every = 2
    slow_target_update: int = 1

    disclam = 0.95
    action_noise = 0.3

    action_size = 0
    obs_shape = 0
    reward_EMA = True
    imag_gradient = 'reinforce'
    imag_horizon = 15

    rssm_node_size = 200
    dyn_hidden = 512
    dyn_deter = 512
    dyn_stoch = 32
    dyn_rec_depth = 1
    dyn_mean_act = 'none'
    dyn_std_act = 'sigmoid2'
    dyn_min_std = 0.1
    dyn_discrete = 32
    unimix_ratio = 0.01
    initial = 'learned'
    act = 'SiLU'
    units = 512
    norm = True
    eval_state_mean = False
    reward_head: Dict = field(default_factory=lambda:{'layers': 2, 'dist': 'symlog_disc', 'loss_scale': 1.0, 'outscale': 0.0})
    cont_head: Dict = field(default_factory=lambda:{'layers': 2, 'loss_scale': 1.0, 'outscale': 1.0})
    encoder: Dict = field(default_factory=lambda:{
        'mlp_keys': '$^', 
        'cnn_keys': 'image',
        'act': 'SiLU',
        'norm': True,
        'cnn_depth': 32,
        'kernel_size': 4,
        'minres': 4,
        'mlp_layers': 5,
        'mlp_units': 1024,
        'symlog_inputs': True
    })
    decoder: Dict = field(default_factory=lambda:{
        'mlp_keys': '$^',
        'cnn_keys': 'image',
        'act': 'SiLU',
        'norm': True,
        'cnn_depth': 32,
        'kernel_size': 4,
        'minres': 4,
        'mlp_layers': 5,
        'mlp_units': 1024,
        'cnn_sigmoid': False,
        'image_dist': 'mse',
        'vector_dist': 'symlog_mse',
        'outscale': 1.0
    })
    actor: Dict = field(default_factory=lambda:{
        'layers': 2,
        'dist': 'normal',
        'entropy': 3e-4,
        'unimix_ratio': 0.01,
        # 'std': 'learned',
        'min_std': 0.1,
        'max_std': 1.0,
        'temp': 0.1,
        'lr': 3e-5,
        'eps': 1e-5,
        'grad_clip': 100.0,
        'outscale': 1.0,
        'dist': 'onehot', 
        'std': 'none'
    })
    critic: Dict = field(default_factory=lambda:{
        'layers': 2,
        'dist': 'symlog_disc',
        'slow_target': True,
        'slow_target_update': 1,
        'slow_target_fraction': 0.02,
        'lr': 3e-5,
        'eps': 1e-5,
        'grad_clip': 100.0,
        'outscale': 0.0
    })
    # scales = dict(
    #         reward=reward_head["loss_scale"],
    #         cont=cont_head["loss_scale"],
    #     )

    def get(self, key, default):
        try:
            return getattr(self, key)
        except AttributeError:
            print(f"警告： 配置中不存在属性'{key}'")
            for line in traceback.format_stack()[:-1]:  # 排除当前函数
                print(line.strip())
            
            return default


